# -*- coding: binary -*-
module Msf

###
#
# This modules provides a target-aware brute-forcing wrapper.  It implements
# the exploit method and calls exploit_brute with target supplied information.
# If the selected target is not a bruteforce target, then single_exploit is
# called.
#
###
module Exploit::Brute

  #
  # Initializes an instance of an exploit module that supports brute force
  # targets.
  #
  def initialize(info = {})
    super

    #
    # Register BruteWait and BruteStep as two advanced options for this
    # exploit even though not all targets may be brute force targets.
    #
    register_advanced_options(
      [
        OptInt.new('BruteWait', [ false, "Delay between brute force attempts"     ]),
        OptInt.new('BruteStep', [ false, "Step size between brute force attempts" ])
      ], Msf::Exploit::Brute)
  end

  #
  # Entry point for initiating an exploit.  This module wrappers the exploit
  # method and determines whether or not the selected target supports brute
  # force.  If it does, it does some special things and wraps the brute
  # forcing logic.
  #
  def exploit
    # Is the selected target a brute force target?
    if (target.bruteforce?)
      # The step direction is automatically calculated
      direction = {}

      bf = target.bruteforce

      # Get the start and stop address hashes
      start = bf.start_addresses ? bf.start_addresses.dup : {}
      stop  = bf.stop_addresses ? bf.stop_addresses.dup : {}
      step  = bf.step_size
      delay = bf.delay

      # Enumerate each start address and try to figure out the direction
      start.each_pair { |name, addr|
        # If there's a stop address, figure out if it's above or below
        # the start address
        if (stop[name])
          if (stop[name] < addr)
            direction[name] = -1
          else
            direction[name] = 1
          end
        # If there's no stop address, infer the direction based on
        # the default
        else
          direction[name] = bf.default_direction
        end
      }

      # Import start/stop address overrides from the datastore
      import_from_datastore(start, 'Start')
      import_from_datastore(stop,  'Stop')

      # User-defined brute wait?
      if self.datastore['BruteWait'] and self.datastore['BruteWait'] > 0
        delay = self.datastore['BruteWait'].to_i
      end

      # User-defined brute step?
      if self.datastore['BruteStep'] and self.datastore['BruteStep'] > 0
        step = self.datastore['BruteStep'].to_i
      end

      # Sane defaults
      delay = 1 if delay.nil? or delay == 0

      # Okay, we've got all this crap out of the way, let's actually brute
      # force
      stopped = []
      curr    = start.dup

      # Automatically determine the step size based off the nop sled length
      if step == 0
        step = payload.nop_sled_size

        if step == 0
          raise OptionValidateError.new(['BruteStep']), "The step size for this exploit is invalid"
        end
      end

      # Keep going until we run out of options
      while (curr.length != stopped.length)

        # Stop brute forcing once a session is found
        break if session_created?

        # Fire off an exploit attempt with the supplied addresses
        brute_exploit(curr)

        # Give it time before we try again
        brute_wait(delay)

        # Scan each current key, increasing it or decreasing it by the
        # step size according to its direction
        curr.each_key { |k|

          # Has movement been stopped on this address?  If so, skip it.
          next if (stopped.include?(k))

          # Calculate the next address before we move it to see if
          # we're going to go over
          next_addr = step * direction[k]

          # If this item has hit a stop address, add it to the stopped
          # hash and move it no further
          if (stop[k])
            if ((direction[k] == 1 and curr[k] + next_addr >= stop[k]) or
                (direction[k] == -1 and curr[k] + next_addr < stop[k]))
              stopped << k
              next
            end
          end

          # If it's not time to stop, move it
          curr[k] += next_addr
        }
      end
    else
      single_exploit
    end
  end

  #
  # This routine is called once per brute force iteration.  The addresses
  # parameter is a hash of addresses that are incremented each iteration and
  # are derived from the target's bruteforce information or the module's
  # datastore in case they are being overriden.
  #
  def brute_exploit(addrs)
  end

  #
  # Call if the target is not a brute force target.
  #
  def single_exploit
  end

  #
  # Waits for the provide delay.
  #
  def brute_wait(delay)
    sleep(delay)
  end

protected

  #
  # Imports information into the supplied hash from the datastore.
  # This is a way of allowing the user to override values for a
  # specific brute force target by name without them actually
  # being conveyed in the options list.  This is a bit of a change
  # from 2.x, but 2.x didn't have per-target brute force
  # addresses, which I think is more valuable.
  #
  def import_from_datastore(hash, prefix = '')
    hash.each_key { |k|
      if (self.datastore[prefix + k])
        hash[k] = self.datastore[prefix + k]
      end
    }
  end

end

end
